module matrix_operations
  use constants
  implicit none
  
contains

  ! 矩阵乘法：c = a * b
  subroutine mtxmtp(n, a, b, c)
    integer, intent(in) :: n
    complex(kind=8), intent(in) :: a(4,4), b(4,4)
    complex(kind=8), intent(out) :: c(4,4)
    
    complex(kind=8) :: zero
    integer :: i, j, k
    
    zero = (0.0d0, 0.0d0)
    c = zero
    
    do j = 1, n
      do k = 1, n
        do i = 1, n
          c(i, j) = c(i,j) + a(i, k)*b(k, j)
        end do
      end do
    end do
  end subroutine mtxmtp
  
  ! 2x2矩阵求逆
  subroutine mtxinv_2(a0)
    real(kind=8), parameter :: epsi = 1.0E-20
    complex(kind=8), intent(inout) :: a0(4,4)
    complex(kind=8) :: a(4,4), d
    
    a(1,1) = a0(1,1)
    a(2,1) = a0(2,1)
    a(1,2) = a0(1,2)
    a(2,2) = a0(2,2)
    
    d = a(1,1)*a(2,2) - a(1,2)*a(2,1)
    if (abs(d) < epsi) then
      print*, 'Singular Matrix_2'
      d = epsi
      stop
    endif
    
    a0(1,1) = a(2,2)/d
    a0(2,2) = a(1,1)/d
    a0(1,2) = -a(1,2)/d
    a0(2,1) = -a(2,1)/d
  end subroutine mtxinv_2
  
  ! 4x4矩阵求逆
  subroutine mtxinv_4(a)
    real(kind=8), parameter :: epsi = 1.0e-10
    complex(kind=8), intent(inout) :: a(4,4)
    real(kind=8) :: d
    integer :: i, j, k, l(4), m(4)
    complex(kind=8) :: t
    
    do k = 1, 4
      d = 0.0
      do i = k, 4
        do j = k, 4
          if (abs(a(i,j)) > d) then
            d = abs(a(i,j))
            l(k) = i
            m(k) = j
          end if
        end do
      end do
      
      if (d < epsi) then
        print*, d
        print*, 'The matrix is singular!'
        stop
      end if
      
      if (l(k) /= k) then
        do j = 1, 4
          t = a(k,j)
          a(k,j) = a(l(k),j)
          a(l(k),j) = t
        end do
      end if
      
      if (m(k) /= k) then
        do i = 1, 4
          t = a(i,k)
          a(i,k) = a(i,m(k))
          a(i,m(k)) = t
        end do
      end if
      
      a(k,k) = 1.0d0 / a(k,k)
      
      do i = 1, 4
        if (i /= k) then
          a(i,k) = -a(k,k) * a(i,k)
        end if
      end do
      
      do i = 1, 4
        do j = 1, 4
          if (i /= k .and. j /= k) then
            a(i,j) = a(i,j) + a(i,k) * a(k,j)
          end if
        end do
      end do
      
      do j = 1, 4
        if (j /= k) then
          a(k,j) = a(k,j) * a(k,k)
        end if
      end do
    end do
    
    do k = 4, 1, -1
      do j = 1, 4
        t = a(k,j)
        a(k,j) = a(m(k),j)
        a(m(k),j) = t
      end do
      
      do i = 1, 4
        t = a(i,k)
        a(i,k) = a(i,l(k))
        a(i,l(k)) = t
      end do
    end do
  end subroutine mtxinv_4
  
  ! 计算第lay层的E矩阵（P-SV情况）
  subroutine mtxe(lay, kn, o, e)
    use parameters, only : media, cpn, csn
    integer, intent(in) :: lay
    real(kind=8), intent(in) :: kn
    complex(kind=8), intent(in) :: o
    complex(kind=8), intent(out) :: e(4,4)
    complex(kind=8) :: c
    
    c = kn*kn + csn(lay)*csn(lay)
    
    e(1,1) = media%vp(lay)*kn/o
    e(1,2) = media%vs(lay)*csn(lay)/o
    e(1,3) = e(1,1)
    e(1,4) = e(1,2)
    e(2,1) = media%vp(lay)*cpn(lay)/o
    e(2,2) = media%vs(lay)*kn/o
    e(2,3) = -e(2,1)
    e(2,4) = -e(2,2)
    e(3,1) = -2*media%vp(lay)*media%mu(lay)*kn*cpn(lay)/o
    e(3,2) = -media%vs(lay)*media%mu(lay)*c/o
    e(3,3) = -e(3,1)
    e(3,4) = -e(3,2)
    e(4,1) = -media%vp(lay)*media%mu(lay)*c/o
    e(4,2) = -2*media%vs(lay)*media%mu(lay)*kn*csn(lay)/o
    e(4,3) = e(4,1)
    e(4,4) = e(4,2)
  end subroutine mtxe
  
end module matrix_operations
